import tensorflow as tf
from tensorflow.keras.layers import SimpleRNN, LSTM, GRU,Dense
from tensorflow.keras.datasets import imdb
from tensorflow.keras.preprocessing import sequence
from tensorflow.keras.models import Sequential
import matplotlib.pyplot as plt
import numpy as np

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  
plt.rcParams['axes.unicode_minus'] = False 

# LSTM单元详细实现
class LSTMCell:
    def __init__(self, input_size, hidden_size):
        self.input_size = input_size
        self.hidden_size = hidden_size
        
        # 权重初始化
        self.W_xi = np.random.randn(hidden_size, input_size) * 0.01
        self.W_hi = np.random.randn(hidden_size, hidden_size) * 0.01
        self.b_i = np.zeros((hidden_size, 1))
        
        self.W_xf = np.random.randn(hidden_size, input_size) * 0.01
        self.W_hf = np.random.randn(hidden_size, hidden_size) * 0.01
        self.b_f = np.zeros((hidden_size, 1))
        
        self.W_xo = np.random.randn(hidden_size, input_size) * 0.01
        self.W_ho = np.random.randn(hidden_size, hidden_size) * 0.01
        self.b_o = np.zeros((hidden_size, 1))
        
        self.W_xc = np.random.randn(hidden_size, input_size) * 0.01
        self.W_hc = np.random.randn(hidden_size, hidden_size) * 0.01
        self.b_c = np.zeros((hidden_size, 1))
    
    def forward(self, x, h_prev, c_prev):
        # 输入门
        i_t = sigmoid(np.dot(self.W_xi, x) + np.dot(self.W_hi, h_prev) + self.b_i)
        
        # 遗忘门
        f_t = sigmoid(np.dot(self.W_xf, x) + np.dot(self.W_hf, h_prev) + self.b_f)
        
        # 输出门
        o_t = sigmoid(np.dot(self.W_xo, x) + np.dot(self.W_ho, h_prev) + self.b_o)
        
        # 候选细胞状态
        c_hat_t = np.tanh(np.dot(self.W_xc, x) + np.dot(self.W_hc, h_prev) + self.b_c)
        
        # 更新细胞状态
        c_t = f_t * c_prev + i_t * c_hat_t
        
        # 更新隐藏状态
        h_t = o_t * np.tanh(c_t)
        
        return h_t, c_t

# LSTM门控机制可视化
def visualize_lstm_gates():
    time_steps = 20
    # 模拟输入序列
    sequence = np.sin(np.linspace(0, 4*np.pi, time_steps))
    
    plt.figure(figsize=(12, 8))
    plt.plot(sequence, 'b-', linewidth=2, label='输入序列')
    plt.title('LSTM门控机制示意图')
    plt.xlabel('时间步')
    plt.ylabel('值')
    plt.legend()
    plt.grid(True)
    plt.show()

visualize_lstm_gates()